package com.lyl.myaiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.*;
import java.util.ArrayList;
import java.util.List;

/**
 * @author lyl
 * @version 1.0
 * @CreateTime: 2025-08-19
 */
public class FileBaseChatMemory implements ChatMemory {

    private final String BASE_DIR;

    private static final Kryo kryo = new Kryo();

    static {
        kryo.setRegistrationRequired(false);
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public FileBaseChatMemory(String dir){
        this.BASE_DIR = dir;
        File baseDir = new File(dir);
        if(!baseDir.exists()){
            baseDir.mkdirs();
        }
    }
    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> messageList = getOrCreatConversation(conversationId);
        messageList.addAll(messages);
        saveConversation(conversationId,messageList);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> messageList = getOrCreatConversation(conversationId);
        return messageList.stream()
                .skip(Math.max(0,(messageList.size() - lastN)))
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        File file = getConversationFile(conversationId);
        if(file.exists()){
            file.delete();
        }
    }

    /**
     * 根据会话id创建文件对象
     * @param conversationId 会话id
     * @return 文件对象
     */
    private File getConversationFile(String conversationId){
        return new File(BASE_DIR,conversationId+".kryo");
    }

    /**
     * 根据会话id获取历史会话记录，如果不存在，返回一个空列表
     * @param conversationId 会话id
     * @return list集合
     */
    private List<Message> getOrCreatConversation(String conversationId){
        File file = getConversationFile(conversationId);
        List<Message> messageList = new ArrayList<>();
        if(file.exists()){
            try(Input input = new Input(new FileInputStream(file))){
                messageList = kryo.readObject(input, ArrayList.class);
            }catch (IOException e){
                e.printStackTrace();
            }
        }
        return messageList;
    }

    private void saveConversation(String conversationId, List<Message> messageList){
        File file = getConversationFile(conversationId);
        try(Output output = new Output(new FileOutputStream(file))){
            kryo.writeObject(output,messageList);
        }catch (IOException e){
            e.printStackTrace();
        }
    }
}
